import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import cv2
import numpy as np

import utils.data_utils as data_utils


def get_layer(model, layer_name):
    """
    Return the targeted layer (nn.Module Object) given a hierarchical layer name,
    separated by /.
    Args:
        model (model): model to get layers from.
        layer_name (str): name of the layer.
    Returns:
        prev_module (nn.Module): the layer from the model with `layer_name` name.
    """
    layer_ls = layer_name.split("/")
    prev_module = model
    for layer in layer_ls:
        prev_module = prev_module._modules[layer]

    print(prev_module)

    return prev_module

def reshape_transform(tensor):
    # (B, C, T, O)
    result = tensor[:, :-1, 1:, :].transpose(2, 3).transpose(1, 2)
    return result

class GradCAM:
    """
    GradCAM class helps create localization maps using the Grad-CAM method for input videos
    and overlap the maps over the input videos as heatmaps.
    https://arxiv.org/pdf/1610.02391.pdf
    """

    def __init__(
        self, model, target_layers, logit_name, colormap="jet"
    ):
        """
        Args:
            model (model): the model to be used.
            target_layers (list of str(s)): name of convolutional layer to be used to get
                gradients and feature maps from for creating localization maps.
            data_mean (tensor or list): mean value to add to input videos.
            data_std (tensor or list): std to multiply for input videos.
            colormap (Optional[str]): matplotlib colormap used to create heatmap.
                See https://matplotlib.org/3.3.0/tutorials/colors/colormaps.html
        """

        self.model = model
        # Run in eval mode.
        self.model.eval()
        self.target_layers = target_layers
        self.logit_name = logit_name

        self.gradients = {}
        self.activations = {}
        self.colormap = plt.get_cmap(colormap)

        self._register_hooks()

    def _register_single_hook(self, layer_name):
        """
        Register forward and backward hook to a layer, given layer_name,
        to obtain gradients and activations.
        Args:
            layer_name (str): name of the layer.
        """

        def get_gradients(module, grad_input, grad_output):
            self.gradients[layer_name] = grad_output[0].detach()

        def get_activations(module, input, output):
            self.activations[layer_name] = output.clone().detach()

        target_layer = get_layer(self.model, layer_name=layer_name)
        target_layer.register_forward_hook(get_activations)
        target_layer.register_full_backward_hook(get_gradients)

    def _register_hooks(self):
        """
        Register hooks to layers in `self.target_layers`.
        """
        for layer_name in self.target_layers:
            self._register_single_hook(layer_name=layer_name)

    def _calculate_localization_map(self, inputs, labels=None):
        """
        Calculate localization map for all inputs with Grad-CAM.
        Args:
            inputs (list of tensor(s)): the input clips.
            labels (Optional[tensor]): labels of the current input clips.
        Returns:
            localization_maps (list of ndarray(s)): the localization map for
                each corresponding input.
            preds (tensor): shape (n_instances, n_class). Model predictions for `inputs`.
        """
        assert len(inputs) == len(
            self.target_layers
        ), "Must register the same number of target layers as the number of input pathways."

        inputs_clone = [{ key: inp[key].clone() if type(inp[key])==torch.Tensor else inp[key].copy() for key in inp.keys() } 
                        for inp in inputs ] 
        preds = self.model(inputs_clone[0])[self.logit_name] # fit model output  preds => (B, NC)


        if labels is None:
            score = torch.max(preds, dim=-1)[0]
        else:  # now only for single label classfication
            if labels.ndim == 1:
                labels = labels.unsqueeze(-1)
            score = torch.gather(preds, dim=1, index=labels)

        self.model.zero_grad()
        score = torch.sum(score)
        score.backward()
        localization_maps = []

        for i, inp in enumerate(inputs):
            if self.logit_name == 'resnet3d':
                _, _, T, H, W = inp['video_frames'].size() 

                gradients = self.gradients[self.target_layers[i]]
                activations = self.activations[self.target_layers[i]]
                B, C, Tg, _, _ = gradients.size()

                weights = torch.mean(gradients.view(B, C, Tg, -1), dim=3)

                weights = weights.view(B, C, Tg, 1, 1)
                localization_map = torch.sum(
                    weights * activations, dim=1, keepdim=True
                )
                localization_map = F.relu(localization_map)
                localization_map = F.interpolate(
                    localization_map,
                    size=(T, H, W),
                    mode="trilinear",
                    align_corners=False,
                )
                localization_map_min, localization_map_max = (
                    torch.min(localization_map.view(B, -1), dim=-1, keepdim=True)[
                        0
                    ],
                    torch.max(localization_map.view(B, -1), dim=-1, keepdim=True)[
                        0
                    ],
                )
                localization_map_min = torch.reshape(
                    localization_map_min, shape=(B, 1, 1, 1, 1)
                )
                localization_map_max = torch.reshape(
                    localization_map_max, shape=(B, 1, 1, 1, 1)
                )
                # Normalize the localization map.
                localization_map = (localization_map - localization_map_min) / (
                    localization_map_max - localization_map_min + 1e-6
                )
                localization_map = localization_map.data
            
            elif self.logit_name == 'stlt':
                _, T, O, _ = inp['boxes'].size()

                gradients = self.gradients[self.target_layers[i]]
                activations = self.activations[self.target_layers[i]]
                _, _, C = gradients.size()

                gradients = gradients.transpose(0, 1)  # BXT, O, C
                gradients = gradients.view(-1, T, O, C)
                gradients = reshape_transform(gradients)

                activations = activations.transpose(0, 1)
                activations = activations.view(-1, T, O, C)
                activations = reshape_transform(activations)
                
                B, C, Tg, _ = gradients.size()

                weights = torch.mean(gradients.view(B, C, Tg, -1), dim=3)
                weights = weights.view(B, C, Tg, 1)
                localization_map = torch.sum(
                    weights * activations, dim=1, keepdim=True
                )

                localization_map_min, localization_map_max = (
                    torch.min(localization_map.view(B, -1), dim=-1, keepdim=True)[
                        0
                    ],
                    torch.max(localization_map.view(B, -1), dim=-1, keepdim=True)[
                        0
                    ],
                )
                localization_map_min = torch.reshape(
                    localization_map_min, shape=(B, 1, 1, 1)
                )
                localization_map_max = torch.reshape(
                    localization_map_max, shape=(B, 1, 1, 1)
                )
                # Normalize the localization map.
                localization_map = (localization_map - localization_map_min) / (
                    localization_map_max - localization_map_min + 1e-6
                )
                localization_map = localization_map.data

            localization_maps.append(localization_map)

        return localization_maps, preds


    def __call__(self, inputs, labels=None):
        """
        Visualize the localization maps on their corresponding inputs as heatmap,
        using Grad-CAM.
        Args:
            inputs (list of tensor(s)): the input clips.
            labels (Optional[tensor]): labels of the current input clips.
            alpha (float): transparency level of the heatmap, in the range [0, 1].
        Returns:
            result_ls (list of tensor(s)): the visualized inputs.
            preds (tensor): shape (n_instances, n_class). Model predictions for `inputs`.
        """
        heatmap_ls = []
        localization_maps, preds = self._calculate_localization_map(
            inputs, labels=labels
        )
        for localization_map in localization_maps:
            # Convert (B, 1, T, H, W) to (B, T, H, W)
            localization_map = localization_map.squeeze(dim=1)
            if localization_map.device != torch.device("cpu"):
                localization_map = localization_map.cpu()
            heatmap = self.colormap(localization_map.numpy())
            if self.logit_name == 'resnet3d':
                heatmap = heatmap[:, :, :, :, :3]
            elif self.logit_name == 'stlt':
                heatmap = heatmap[:, :, :, :3]

            heatmap_ls.append(heatmap)

        return heatmap_ls, preds


def show_cam_on_image(rgb_imgs, heatmaps, data_mean=[0.5, 0.5, 0.5], data_std=[0.5, 0.5, 0.5], alpha=0.5):

    result_ls = []
    for i, heatmap in enumerate(heatmaps):
        # Permute input from (B, C, T, H, W) to (B, T, H, W, C)
        curr_inp = rgb_imgs[i].permute(0, 2, 3, 4, 1)
        if curr_inp.device != torch.device("cpu"):
            curr_inp = curr_inp.cpu()
        curr_inp = data_utils.revert_tensor_normalize(
            curr_inp, data_mean, data_std
        )
        heatmap = torch.from_numpy(heatmap)
        curr_inp = alpha * heatmap + (1 - alpha) * curr_inp
        # Permute inp to (B, T, C, H, W)
        curr_inp = curr_inp.permute(0, 1, 4, 2, 3)
        result_ls.append(curr_inp)
    
    return result_ls


def show_box_cam_on_image(rgb_imgs, heatmaps, boxes, alpha=0.5):

    result_ls = []
    boxes = boxes[:, :-1, 1:, :]  # B x T x O x 4
    for i, heatmap in enumerate(heatmaps):
        # Permute input from (B, C, T, H, W) to (B, T, H, W, C)
        curr_inp = rgb_imgs[i].permute(0, 2, 3, 4, 1)
        if curr_inp.device != torch.device("cpu"):
            curr_inp = curr_inp.cpu()
        if boxes.device != torch.device("cpu"):
            boxes = boxes.cpu()

        full_heatmap = torch.zeros_like(curr_inp)
        _, _, H, W, _ = full_heatmap.size()
        video_size = torch.tensor([W, H]).repeat(2)
        boxes = (boxes * video_size).to(torch.int64)  # box => (x1, y1, x2, y2)

        heatmap = torch.from_numpy(heatmap)

        B, T, O, _ = heatmap.size()
        for b in range(B):
            for t in range(T):
                for o in range(O):
                    x1 = boxes[b, t, o, 0]
                    y1 = boxes[b, t, o, 1]
                    x2 = boxes[b, t, o, 2]
                    y2 = boxes[b, t, o, 3]

                    if x1==0 and x2==0 and y1==0 and y2==0: 
                        continue
                    full_heatmap[b, t, y1:y2+1, x1:x2+1, :] = heatmap[b, t, o, :]
        
        curr_inp = alpha * full_heatmap + (1 - alpha) * curr_inp
        # Permute inp to (B, T, C, H, W)
        curr_inp = curr_inp.permute(0, 1, 4, 2, 3)
        result_ls.append(curr_inp)

    return result_ls



        